#include <immintrin.h>
#include <cassert>
#include <cstring> // for memcpy
#include <tuple>

/**
 * AVX2优化的2D卷积函数（最终完整版本）
 * @param dst        输出图像数据（内存布局：dst_stride行距）
 * @param src        输入图像数据（内存布局：src_stride行距）
 * @param width      图像宽度（像素）
 * @param height     图像高度（像素）
 * @param kernel     3x3卷积核
 * @param src_stride 输入图像行距（float数）
 * @param dst_stride 输出图像行距（float数）
 *
 * 优化特性：
 * 1. 行数据复用：prev_row复用上一轮的curr_row
 * 2. 滑动窗口：通过跨通道位移减少内存加载
 * 3. 智能预取：按缓存行预取数据
 * 4. 边界保护：安全加载防止越界
 */
void filter2D_avx2_final_optimized(float* dst,
                                   const float* src,
                                   int width,
                                   int height,
                                   const float (&kernel)[3][3],
                                   int src_stride,
                                   int dst_stride)
{
    constexpr int VF = 8; // AVX2向量包含8个float
    const int main_width = width - 2; // 考虑边界裁剪
    const int prologue = 1; // 边界处理
    const int aligned_width = ((main_width - prologue) / VF) * VF + prologue;

    // 加载卷积核到寄存器
    const __m256 k0 = _mm256_set1_ps(kernel[0][0]);//load fix data
    const __m256 k1 = _mm256_set1_ps(kernel[0][1]);
    const __m256 k2 = _mm256_set1_ps(kernel[0][2]);
    const __m256 k3 = _mm256_set1_ps(kernel[1][0]);
    const __m256 k4 = _mm256_set1_ps(kernel[1][1]);
    const __m256 k5 = _mm256_set1_ps(kernel[1][2]);
    const __m256 k6 = _mm256_set1_ps(kernel[2][0]);
    const __m256 k7 = _mm256_set1_ps(kernel[2][1]);
    const __m256 k8 = _mm256_set1_ps(kernel[2][2]);

    // 初始化三行指针
    const float* prev_row = src + 0 * src_stride; // y-1
    const float* curr_row = src + 1 * src_stride; // y
    const float* next_row = src + 2 * src_stride; // y+1

    // 列方向数据复用寄存器
    __m256 prev0, prev1, curr0, curr1, next0, next1;

    // 添加边界指针
    const float* src_end = src + (height-1)*src_stride + width;

    // 安全加载函数
    auto safe_load = [](const float* addr, const float* end) -> __m256 {
        if (addr + 7 >= end) {
            float tmp[8] = {0};
            int valid = end - addr;
            if (valid > 0) memcpy(tmp, addr, valid*sizeof(float));
            return _mm256_loadu_ps(tmp);//load 非对齐的数据
        }
        return _mm256_loadu_ps(addr);
    };

    // 窗口生成函数
    auto generate_window = [](__m256 low, __m256 high) -> std::tuple<__m256, __m256, __m256> {
        // 合并高低通道
        __m256 merged_hi = _mm256_permute2f128_ps(low, high, 0x21);
        
        // 生成x位置窗口
        __m256 shifted1 = _mm256_permute_ps(low, _MM_SHUFFLE(2,1,0,3));
        __m256 x1 = _mm256_blend_ps(shifted1, merged_hi, 0x11);
        
        // 生成x+1位置窗口
        __m256 shifted2 = _mm256_permute_ps(low, _MM_SHUFFLE(1,0,3,2));
        __m256 x2 = _mm256_blend_ps(shifted2, merged_hi, 0x33);
        
        return std::make_tuple(low, x1, x2);
    };

    for (int y = 1; y < height-2; ++y) { // 修改循环边界
        float* dst_row = dst + y*dst_stride;

        // 保存当前行数据指针（用于下次循环）
        const float* saved_curr = curr_row;
        const float* saved_next = next_row;

        // 主循环处理
        for (int x = 1; x < aligned_width; x += VF) {
            // 数据加载策略：
            // 1. 第一列：加载所有数据
            // 2. 后续列：复用前一次加载的高位数据
            if (x == 1) {
                // 首次加载：完整加载三行数据
                prev0 = safe_load(prev_row + x-1, src_end);
                prev1 = safe_load(prev_row + x+7, src_end);
                curr0 = safe_load(curr_row + x-1, src_end);
                curr1 = safe_load(curr_row + x+7, src_end);
                next0 = safe_load(next_row + x-1, src_end);
                next1 = safe_load(next_row + x+7, src_end);
            } else {
                // 复用前一次加载的高位数据
                prev0 = prev1;
                curr0 = curr1;
                next0 = next1;
                
                // 仅加载新的高位数据
                prev1 = safe_load(prev_row + x+7, src_end);
                curr1 = safe_load(curr_row + x+7, src_end);
                next1 = safe_load(next_row + x+7, src_end);
            }

            // 生成三行窗口
            __m256 p0, p1, p2;
            std::tie(p0, p1, p2) = generate_window(prev0, prev1);
            __m256 prev_x0 = p0, prev_x1 = p1, prev_x2 = p2;

            std::tie(p0, p1, p2) = generate_window(curr0, curr1);
            __m256 curr_x0 = p0, curr_x1 = p1, curr_x2 = p2;

            std::tie(p0, p1, p2) = generate_window(next0, next1);
            __m256 next_x0 = p0, next_x1 = p1, next_x2 = p2;

            // 卷积计算
            __m256 sum = _mm256_fmadd_ps(prev_x0, k0, _mm256_setzero_ps());
            sum = _mm256_fmadd_ps(prev_x1, k1, sum);
            sum = _mm256_fmadd_ps(prev_x2, k2, sum);
            sum = _mm256_fmadd_ps(curr_x0, k3, sum);
            sum = _mm256_fmadd_ps(curr_x1, k4, sum);
            sum = _mm256_fmadd_ps(curr_x2, k5, sum);
            sum = _mm256_fmadd_ps(next_x0, k6, sum);
            sum = _mm256_fmadd_ps(next_x1, k7, sum);
            sum = _mm256_fmadd_ps(next_x2, k8, sum);

            // 边界掩码处理
            if (x + VF > main_width) {
                __m256 mask = _mm256_cmp_ps(_mm256_set1_ps(x + 0.0f), 
                    _mm256_set1_ps(main_width-1.0f), _CMP_LT_OQ);
                sum = _mm256_blendv_ps(_mm256_setzero_ps(), sum, mask);
            }

            // 存储结果
            _mm256_storeu_ps(dst_row + x, sum);

            // 预取下一组数据（提前4个向量）
            if (x % (VF*4) == 0) {
                _mm_prefetch((const char*)(prev_row + x + VF*4), _MM_HINT_T0);
                _mm_prefetch((const char*)(curr_row + x + VF*4), _MM_HINT_T0);
                _mm_prefetch((const char*)(next_row + x + VF*4), _MM_HINT_T0);
            }
        }

        // 处理剩余元素
        for (int x = aligned_width; x < main_width; ++x) {
            float sum = 0;
            for (int ky = -1; ky <= 1; ++ky) {
                for (int kx = -1; kx <= 1; ++kx) {
                    sum += curr_row[x + kx + ky*src_stride] * 
                           kernel[ky+1][kx+1];
                }
            }
            dst_row[x] = sum;
        }

        // 更新行指针（关键优化！）
        prev_row = saved_curr; // 下次循环的prev_row = 当前curr_row
        curr_row = saved_next; // 下次循环的curr_row = 当前next_row
        next_row = src + (y+2)*src_stride; // 加载新行
    }

    // 单独处理最后两行
    for (int y = height-2; y < height-1; ++y) {
        float* dst_row = dst + y*dst_stride;
        const float* prev_row = src + (y-1)*src_stride;
        const float* curr_row = src + y*src_stride;

        for (int x = 1; x < main_width; ++x) {
            float sum = 0;
            for (int ky = -1; ky <= 0; ++ky) { // 仅处理前两行
                for (int kx = -1; kx <= 1; ++kx) {
                    sum += curr_row[x + kx + ky*src_stride] * 
                           kernel[ky+1][kx+1];
                }
            }
            dst_row[x] = sum;
        }
    }
}
